import fitz  # PyMuPDF
from typing import List, Dict, Tuple, Optional, Any
import re

class AdvancedPDFTableMerger:
    def __init__(self, pdf_path: str):
        self.doc = fitz.open(pdf_path)
        self.pages = [self.doc.load_page(i) for i in range(len(self.doc))]

    def extract_page_elements(self, page_num: int) -> Dict[str, Any]:
        """
        提取页面中的所有元素：文本块、线条和Clip区域
        """
        page = self.pages[page_num]
        elements = {
            "text_blocks": [],
            "lines": [],
            "clip_regions": [],
            "page_rect": page.rect
        }

        # 提取文本块
        text_blocks = page.get_text("dict")["blocks"]
        for block in text_blocks:
            if block["type"] == 0:  # 文本块
                elements["text_blocks"].append({
                    "bbox": block["bbox"],
                    "text": "\n".join([" ".join([span["text"] for span in line["spans"]])
                                      for line in block.get("lines", [])])
                })

        # 提取线条
        drawings = page.get_drawings()
        for drawing in drawings:
            if drawing["items"]:  # 检查是否有绘制项
                for item in drawing["items"]:
                    if item[0] == "l":  # 线条
                        # x0, y0, x1, y1 = item[1:]
                        x0, y0, x1, y1 = item[1]
                        elements["lines"].append({
                            "type": "line",
                            "points": [(x0, y0), (x1, y1)],
                            "bbox": (min(x0, x1), min(y0, y1), max(x0, x1), max(y0, y1))
                        })
                    elif item[0] == "re":  # 矩形
                        print(item)
                        # x, y, w, h = item[1:]
                        x, y, w, h = item[1]
                        elements["lines"].append({
                            "type": "rect",
                            "bbox": (x, y, x + w, y + h)
                        })

            # 提取Clip区域
            if "clip" in drawing and drawing["clip"]:
                clip_rect = drawing["rect"]
                elements["clip_regions"].append(clip_rect)

        return elements

    def find_table_candidates(self, page_num: int) -> List[Dict[str, Any]]:
        """
        在页面中查找可能的表格区域
        基于文本对齐和线条特征
        """
        elements = self.extract_page_elements(page_num)
        text_blocks = elements["text_blocks"]
        lines = elements["lines"]
        page_rect = elements["page_rect"]

        # 按Y坐标分组文本块，形成潜在的行
        text_blocks.sort(key=lambda x: x["bbox"][1])  # 按Y坐标排序
        rows = []
        current_row = []
        last_y = None

        for block in text_blocks:
            bbox = block["bbox"]
            y_center = (bbox[1] + bbox[3]) / 2

            if last_y is None or abs(y_center - last_y) > 5:  # 新的行
                if current_row:
                    rows.append(current_row)
                current_row = [block]
                last_y = y_center
            else:
                current_row.append(block)

        if current_row:
            rows.append(current_row)

        # 检测表格结构：检查列对齐
        table_candidates = []

        for i, row in enumerate(rows):
            # 按X坐标排序行中的文本块
            row.sort(key=lambda x: x["bbox"][0])

            # 检查是否有多个文本块在一行中（可能表示表格列）
            if len(row) > 1:
                # 计算列边界
                columns = []
                for block in row:
                    columns.append((block["bbox"][0], block["bbox"][2]))

                # 检查是否有线条与这些列对齐
                aligned_lines = []
                for line in lines:
                    line_bbox = line["bbox"]
                    for col_start, col_end in columns:
                        # 检查线条是否与列对齐
                        if (abs(line_bbox[0] - col_start) < 5 and
                            abs(line_bbox[2] - col_end) < 5):
                            aligned_lines.append(line)

                # 如果有对齐的线条，可能是一个表格行
                if aligned_lines:
                    # 计算行的边界框
                    row_bbox = (
                        min(block["bbox"][0] for block in row),
                        min(block["bbox"][1] for block in row),
                        max(block["bbox"][2] for block in row),
                        max(block["bbox"][3] for block in row)
                    )

                    table_candidates.append({
                        "row_index": i,
                        "bbox": row_bbox,
                        "text_blocks": row,
                        "aligned_lines": aligned_lines
                    })

        return table_candidates

    def detect_cross_page_tables(self) -> List[Tuple[int, int]]:
        """
        检测跨页表格
        基于多种特征：Clip区域、表格线对齐、文本连续性
        """
        cross_page_tables = []

        for i in range(len(self.pages) - 1):
            page1_elements = self.extract_page_elements(i)
            page2_elements = self.extract_page_elements(i + 1)

            # 方法1: 检查Clip区域对齐
            clip_pairs = []
            for clip1 in page1_elements["clip_regions"]:
                for clip2 in page2_elements["clip_regions"]:
                    if self._are_rects_aligned(clip1, clip2, "horizontal", 10):
                        clip_pairs.append((clip1, clip2))

            # 方法2: 检查表格线对齐
            line_pairs = []
            page1_lines = [line for line in page1_elements["lines"]
                          if line["bbox"][3] > page1_elements["page_rect"].y1 - 20]  # 底部20点内的线条
            page2_lines = [line for line in page2_elements["lines"]
                          if line["bbox"][1] < page2_elements["page_rect"].y0 + 20]  # 顶部20点内的线条

            for line1 in page1_lines:
                for line2 in page2_lines:
                    if self._are_lines_aligned(line1, line2, 10):
                        line_pairs.append((line1, line2))

            # 方法3: 检查文本连续性
            text_continuity = self._check_text_continuity(
                page1_elements["text_blocks"],
                page2_elements["text_blocks"]
            )

            # 如果有任何跨页特征，则认为有跨页表格
            if clip_pairs or line_pairs or text_continuity:
                cross_page_tables.append((i, i + 1))
                print(f"检测到跨页表格: 第 {i+1} 页和第 {i+2} 页")
                if clip_pairs:
                    print(f"  - 找到 {len(clip_pairs)} 对对齐的Clip区域")
                if line_pairs:
                    print(f"  - 找到 {len(line_pairs)} 对对齐的线条")
                if text_continuity:
                    print(f"  - 文本连续性得分: {text_continuity}")

        return cross_page_tables

    def _are_rects_aligned(self, rect1, rect2, direction: str, tolerance: float) -> bool:
        """检查两个矩形是否在指定方向上对齐"""
        if direction == "horizontal":
            return (abs(rect1.x0 - rect2.x0) < tolerance and
                    abs(rect1.x1 - rect2.x1) < tolerance)
        elif direction == "vertical":
            return (abs(rect1.y0 - rect2.y0) < tolerance and
                    abs(rect1.y1 - rect2.y1) < tolerance)
        return False

    def _are_lines_aligned(self, line1, line2, tolerance: float) -> bool:
        """检查两条线是否对齐"""
        bbox1 = line1["bbox"]
        bbox2 = line2["bbox"]

        # 检查水平线对齐
        if abs(bbox1[1] - bbox1[3]) < 2 and abs(bbox2[1] - bbox2[3]) < 2:  # 水平线
            return (abs(bbox1[0] - bbox2[0]) < tolerance and
                    abs(bbox1[2] - bbox2[2]) < tolerance)

        # 检查垂直线对齐
        elif abs(bbox1[0] - bbox1[2]) < 2 and abs(bbox2[0] - bbox2[2]) < 2:  # 垂直线
            return (abs(bbox1[1] - bbox2[1]) < tolerance and
                    abs(bbox1[3] - bbox2[3]) < tolerance)

        return False

    def _check_text_continuity(self, text_blocks1, text_blocks2) -> float:
        """检查两页文本之间的连续性"""
        if not text_blocks1 or not text_blocks2:
            return 0.0

        # 获取第一页底部和第二页顶部的文本
        bottom_texts = [block["text"] for block in text_blocks1
                       if block["bbox"][3] > self.pages[0].rect.y1 - 50]  # 底部50点内的文本
        top_texts = [block["text"] for block in text_blocks2
                    if block["bbox"][1] < self.pages[0].rect.y0 + 50]  # 顶部50点内的文本

        if not bottom_texts or not top_texts:
            return 0.0

        # 简单的连续性检查：检查数字序列、日期序列等
        bottom_text = " ".join(bottom_texts)
        top_text = " ".join(top_texts)

        # 检查数字连续性
        bottom_numbers = re.findall(r'\d+', bottom_text)
        top_numbers = re.findall(r'\d+', top_text)

        if bottom_numbers and top_numbers:
            try:
                last_bottom = int(bottom_numbers[-1])
                first_top = int(top_numbers[0])
                if first_top == last_bottom + 1:
                    return 1.0  # 数字连续
            except:
                pass

        # 检查文本相似性（简单版）
        bottom_words = set(re.findall(r'\w+', bottom_text.lower()))
        top_words = set(re.findall(r'\w+', top_text.lower()))
        common_words = bottom_words & top_words

        if common_words:
            return min(0.9, len(common_words) / 10)  # 基于共同词数量的得分

        return 0.0

    def extract_table_data(self, page_num: int, region: Optional[Tuple[float, float, float, float]] = None) -> List[List[str]]:
        """
        提取指定区域内的表格数据
        如果未指定区域，则提取整个页面中的表格数据
        """
        elements = self.extract_page_elements(page_num)
        text_blocks = elements["text_blocks"]

        if region:
            # 只提取指定区域内的文本块
            x0, y0, x1, y1 = region
            text_blocks = [block for block in text_blocks
                          if (x0 <= block["bbox"][0] <= x1 and
                              y0 <= block["bbox"][1] <= y1 and
                              x0 <= block["bbox"][2] <= x1 and
                              y0 <= block["bbox"][3] <= y1)]

        # 按行组织文本
        text_blocks.sort(key=lambda x: x["bbox"][1])  # 按Y坐标排序
        rows = []
        current_row = []
        last_y = None

        for block in text_blocks:
            bbox = block["bbox"]
            y_center = (bbox[1] + bbox[3]) / 2

            if last_y is None or abs(y_center - last_y) > 5:  # 新的行
                if current_row:
                    # 按X坐标排序当前行中的文本块
                    current_row.sort(key=lambda x: x["bbox"][0])
                    rows.append([block["text"] for block in current_row])
                current_row = [block]
                last_y = y_center
            else:
                current_row.append(block)

        if current_row:
            current_row.sort(key=lambda x: x["bbox"][0])
            rows.append([block["text"] for block in current_row])

        return rows

    def merge_cross_page_tables(self):
        """
        合并检测到的跨页表格
        """
        cross_page_pairs = self.detect_cross_page_tables()
        merged_tables = {}

        for page1, page2 in cross_page_pairs:
            print(f"\n处理跨页表格: 第 {page1+1} 页和第 {page2+1} 页")

            # 提取两页的表格数据
            table_part1 = self.extract_table_data(page1)
            table_part2 = self.extract_table_data(page2)

            # 简单合并：将第二页的数据追加到第一页
            merged_table = table_part1 + table_part2

            merged_tables[(page1, page2)] = merged_table
            print(f"合并后的表格有 {len(merged_table)} 行")

            # 显示前几行
            for i, row in enumerate(merged_table[:3]):
                print(f"行 {i+1}: {row}")
            if len(merged_table) > 3:
                print("...")

        return merged_tables

# 使用示例
if __name__ == "__main__":
    import sys
    pdf_path = sys.argv[1]

    try:
        merger = AdvancedPDFTableMerger(pdf_path)
        merged_tables = merger.merge_cross_page_tables()

        if not merged_tables:
            print("未检测到跨页表格")

    except Exception as e:
        print(f"处理PDF时发生错误: {e}")
        import traceback
        traceback.print_exc()
